from fastapi import FastAPI, File, UploadFile, Form
from typing import Annotated
from langchain_helper  import load_pdf_and_save_to_index, load_index, query_index

import  os

file_upload_path='./uploads'

app = FastAPI()

@app.get('/')
def read_root():
    return "hello world"


@app.post('/upload')
def upload(
    file: Annotated[UploadFile, File()],
    index_name: Annotated[str, Form()]
):
    file_target_path = os.path.join(file_upload_path, file.filename)
    with open(file_target_path, 'wb') as f:
        contents = file.file.read()
        f.write(contents)
    load_pdf_and_save_to_index(file_target_path, index_name)
    return {"filename": file.filename, "index_name": index_name}


from pydantic import BaseModel
class Query(BaseModel):
    index_name: str
    query: str


@app.post("/query")
def query(request: Query):
    index_name = request.index_name
    query = request.query
    index = load_index(index_name=index_name)
    ans = query_index(index, query)
    return {"answer": ans, "index_name": index_name, "query": query}

